import pyautogui
import pygetwindow as gw
import time
import os
import sys

# Windows 控制台 UTF-8 设置（避免 emoji 编码错误）
if sys.platform == 'win32':
    try:
        # 设置标准输出为 UTF-8
        import io
        sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace', line_buffering=True)
        sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8', errors='replace', line_buffering=True)
    except Exception:
        # 如果设置失败，静默忽略
        pass

def locate_and_click(image_path, description, confidence=0.85, timeout=6):
    """
    在屏幕上查找指定图片并点击中心点，允许一定的等待时间。
    """
    if not os.path.exists(image_path):
        print(f"❌ 未找到截图文件: {image_path}")
        return False

    start = time.time()
    while time.time() - start < timeout:
        try:
            location = pyautogui.locateCenterOnScreen(image_path)
        except AttributeError:
            # 未安装 opencv-python 时，confidence 参数不可用
            location = pyautogui.locateCenterOnScreen(image_path)
        except pyautogui.ImageNotFoundException:
            location = None

        if location:
            pyautogui.moveTo(location)
            pyautogui.click()
            return True
        time.sleep(0.5)

    print(f"❌ 屏幕上未找到目标: {description}（截图: {image_path}）")
    return False


def export_tdx_block_via_34():
    """
    通过通达信快捷键 '3' -> '4' 打开数据导出，
    自动选择“板块成分”并导出为 CSV。
    """

    # 1. 激活通达信窗口（增强版）
    windows = gw.getWindowsWithTitle('通达信')
    if not windows:
        print("❌ 未找到通达信窗口，请确保软件已打开")
        return False
    
    tdx = windows[0]
    
    # 先恢复窗口（如果是最小化状态）
    if tdx.isMinimized:
        tdx.restore()
        time.sleep(0.3)
    
    # 多次激活确保获得焦点
    for i in range(3):
        tdx.activate()
        time.sleep(0.3)
    
    # 先按ESC退出可能存在的弹窗
    pyautogui.press('esc')
    time.sleep(0.3)
    pyautogui.press('esc')
    time.sleep(0.5)

    # 点击窗口中心确保焦点
    try:
        center_x = tdx.left + tdx.width // 4
        center_y = tdx.top + tdx.height // 4
        pyautogui.click(center_x, center_y)
        time.sleep(0.3)
    except Exception as e:
        print(f"⚠️ 窗口点击失败: {e}")
    
    print("✅ 通达信窗口已激活")


    # 先按ESC退出可能存在的弹窗
    pyautogui.press('esc')
    time.sleep(0.3)
    pyautogui.press('esc')
    time.sleep(0.5)
    
    # 使用更可靠的按键方式：按下后稍微保持
    print("🔑 按下快捷键 '3' 进入分析菜单...")
    pyautogui.keyDown('3')
    time.sleep(0.05)  # 保持按下50毫秒
    pyautogui.keyUp('3')
    time.sleep(1.0)  # 增加等待时间，确保菜单弹出
    
    print("🔑 按下快捷键 '4' 打开数据导出...")
    pyautogui.keyDown('4')
    time.sleep(0.05)
    pyautogui.keyUp('4')
    time.sleep(2.5)  # 增加等待时间，确保导出窗口弹出
    
    print("✅ 快捷键已发送")
    
    time.sleep(0.5)
    
    
    # 3. 导航到“板块成分”（通常第3项：K线=1, 分笔=2, 板块成分=3）
    # 默认选中第1项，按 ↓ 两次到第3项
    pyautogui.press('enter')  # 进入板块成分导出设置
    time.sleep(1.5)

    # 4. 聚焦“开始导出”按钮
    # 经测试，多数版本只需 Tab 一次即可到“开始导出”
    step1_img = "D:\\Quant\\A\\tdxcontrol\\data\\step1.png"
    if not locate_and_click(step1_img, "板块成分导出按钮"):
        return False
    time.sleep(1.5)  # 等待“另存为”对话框

    # 5. 输入文件名并保存
    step2_img = "D:\\Quant\\A\\tdxcontrol\\data\\step2.png"
    if not locate_and_click(step2_img, "保存文件位置"):
        return False
    time.sleep(1.5)  # 等待“另存为”对话框


    # 6. 关闭导出窗口
    pyautogui.press('esc')
    time.sleep(0.3)
    pyautogui.press('esc')
    time.sleep(0.5)
    pyautogui.press('esc')
    time.sleep(0.3)
    pyautogui.press('esc')
    time.sleep(0.5)

    # 点击窗口确保焦点
    try:    
        center_x = tdx.left + tdx.width // 4
        center_y = tdx.top + tdx.height // 4
        pyautogui.click(center_x, center_y)
        time.sleep(0.3)
    except Exception as e:
        print(f"⚠️ 窗口点击失败: {e}")

    # 6. 关闭导出窗口
    pyautogui.press('esc')
    time.sleep(0.3)
    pyautogui.press('esc')
    time.sleep(0.5)
    pyautogui.press('esc')
    time.sleep(0.3)
    pyautogui.press('esc')
    time.sleep(0.5)

    if not tdx.isMinimized:
        tdx.minimize()

    # 验证文件是否生成
    save_path =  f"D:\\Quant\\quantdata\\tdxexport\\行业板块.txt"
    if os.path.exists(save_path):
        print(f"✅ 板块成分已成功导出到: {save_path}")
        return True
    else:
        print(f"⚠️ 文件未生成，可能保存失败: {save_path}")
        return False

    

# 使用示例
if __name__ == "__main__":
    success = export_tdx_block_via_34()
    if not success:
        print("导出失败，请检查通达信窗口状态和版本兼容性。")
